/*
 * Copyright (c)  2022.  Zhengcaiyun Co., Ltd.
 */

package cn.hy.redis.limiter;

import cn.hy.redis.limiter.util.TimeUtils;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import java.util.Collections;
import java.util.concurrent.TimeUnit;
import static com.google.common.base.Preconditions.checkArgument;

@Slf4j
public class RedisLuaRateLimiter implements RateLimiter {

    /**
     * 限流逻辑脚本
     */
    private static final String LIMITER_LOGIC_SCRIPT = "local rateLimiterKey, permitsPerSecond, maxBurstSeconds, needPermits, currentMillis = KEYS[1], tonumber(ARGV[1]), tonumber(ARGV[2]), tonumber(ARGV[3]), tonumber(ARGV[4])\n" +
            "local rateLimiterPermits = redis.call('HGETALL', rateLimiterKey)\n" +
            "\n" +
            "local tempMaxPermits = permitsPerSecond * maxBurstSeconds\n" +
            "\n" +
            "local maxPermits, stableIntervalMillis, storedPermits, nextFreeTicketMillis, returnValue\n" +
            "\n" +
            "if table.maxn(rateLimiterPermits) >= 8 then\n" +
            "    maxPermits, stableIntervalMillis, storedPermits, nextFreeTicketMillis = tonumber(rateLimiterPermits[2]), tonumber(rateLimiterPermits[4]), tonumber(rateLimiterPermits[6]), tonumber(rateLimiterPermits[8])\n" +
            "\n" +
            "    -- The current limiting strategy has changed\n" +
            "    if 0.5 < math.abs(maxPermits - tempMaxPermits) then\n" +
            "        maxPermits, stableIntervalMillis = tempMaxPermits, 1000.00 / permitsPerSecond\n" +
            "        storedPermits, nextFreeTicketMillis = permitsPerSecond, currentMillis\n" +
            "    end\n" +
            "else\n" +
            "    maxPermits, stableIntervalMillis = tempMaxPermits, 1000.00 / permitsPerSecond\n" +
            "    storedPermits, nextFreeTicketMillis, returnValue = permitsPerSecond, currentMillis, currentMillis\n" +
            "end\n" +
            "\n" +
            "-- Current time has exceeded the most recent available time, reset\n" +
            "if currentMillis > nextFreeTicketMillis then\n" +
            "    storedPermits = math.min(maxPermits, storedPermits + (currentMillis - nextFreeTicketMillis) / stableIntervalMillis);\n" +
            "    nextFreeTicketMillis = currentMillis\n" +
            "end\n" +
            "\n" +
            "returnValue = nextFreeTicketMillis\n" +
            "\n" +
            "local storedPermitsToSpend = math.min(needPermits, storedPermits)\n" +
            "local freshPermits = needPermits - storedPermitsToSpend\n" +
            "local waitMillis = math.floor(freshPermits * stableIntervalMillis)\n" +
            "\n" +
            "nextFreeTicketMillis = returnValue + waitMillis\n" +
            "storedPermits = storedPermits - storedPermitsToSpend\n" +
            "\n" +
            "redis.call('HMSET', rateLimiterKey, 'maxPermits', maxPermits, 'stableIntervalMicros', stableIntervalMillis, 'storedPermits', storedPermits, 'nextFreeTicketMillis', nextFreeTicketMillis)\n" +
            "\n" +
            "-- 1 min\n" +
            "redis.call('EXPIRE', rateLimiterKey, 60)\n" +
            "\n" +
            "return math.max(returnValue - currentMillis, 0)";

    /**
     * redis client
     */
    private StringRedisTemplate stringRedisTemplate;

    /**
     * load script
     */
    private final DefaultRedisScript<Long> scriptLoad = new DefaultRedisScript<>(LIMITER_LOGIC_SCRIPT, Long.class);

    /**
     * 获取数据
     */
    private final String lua1 = "local n=redis.call('HGET',KEYS[1],'nextFreeTicketMillis')\n" +
            "if (n==nil or (type(n)=='boolean' and n==false)) then return 1 end\n" +
            "if (n-tonumber(ARGV[1])) <= tonumber(ARGV[2]) then return 1 else return 0 end";
    /**
     * load script
     */
    private final DefaultRedisScript<Long> scriptLoadForGet = new DefaultRedisScript<>(lua1,Long.class);

    /**
     * redis key
     */
    private String rateLimiterKey;

    /**
     * 每秒存入的令牌数
     */
    private double permitsPerSecond;

    /**
     * 最大存储maxBurstSeconds秒生成的令牌
     */
    private double maxBurstSeconds;


    public RedisLuaRateLimiter(StringRedisTemplate stringRedisTemplate, String rateLimiterKey, double permitsPerSecond, double maxBurstSeconds) {
        checkArgument(stringRedisTemplate != null, "The redis client must be initialized first");
        checkArgument(StringUtils.isNotBlank(rateLimiterKey), "The rateLimiterKey must be not blank");
        checkArgument(0.0 < permitsPerSecond, "The rateLimiterKey must be positive");
        checkArgument(0.0 <= maxBurstSeconds, "The maxBurstSeconds must not be negative");
        this.rateLimiterKey = rateLimiterKey;
        this.permitsPerSecond = permitsPerSecond;
        this.maxBurstSeconds = maxBurstSeconds;
        this.stringRedisTemplate = stringRedisTemplate;
    }

    @Override
    public long acquire(int permits) {
        long millisToWait = this.reserve(permits);
        TimeUtils.sleepUninterruptibly(millisToWait, TimeUnit.MILLISECONDS);
        return millisToWait;
    }

    @Override
    public boolean tryAcquire(int permits, long timeout, TimeUnit unit) {
        checkPermits(permits);
        long timeoutMillis = Math.max(unit.toMillis(timeout), 0);
        long millisToWait;
        if (!this.canAcquire(timeoutMillis)) {
            return false;
        } else {
            millisToWait = this.reserveAndGetWaitLength(permits);
        }
        TimeUtils.sleepUninterruptibly(millisToWait, TimeUnit.MILLISECONDS);
        return true;
    }

    private long reserve(int permits) {
        this.checkPermits(permits);
        return this.reserveAndGetWaitLength(permits);
    }

    private void checkPermits(int permits) {
        checkArgument(permits > 0, "Requested permits " + permits + " must be positive");
    }

    private boolean canAcquire(long timeoutMillis) {
        try {
            Long execute = stringRedisTemplate.execute(scriptLoadForGet,
                    Collections.singletonList(rateLimiterKey),
                    String.valueOf(timeoutMillis),
                    String.valueOf(System.currentTimeMillis()));
            return execute != null && execute == 1;
        } catch (Exception ex) {
            throw new RuntimeException(ex.getMessage(), ex.getCause());
        }
    }

    private long reserveAndGetWaitLength(int permits) {
        try {
            Long r = stringRedisTemplate.execute(scriptLoad,
                    Collections.singletonList(rateLimiterKey),
                    String.valueOf(permitsPerSecond),
                    String.valueOf(maxBurstSeconds),
                    String.valueOf(permits),
                    String.valueOf(TimeUtils.currentTimeMillis())
            );
            return r != null ? r : 0;
        } catch (Exception ex) {
            throw new RuntimeException(ex.getMessage(), ex.getCause());
        }
    }
}
